#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@IDE     ：PyCharm 
@Author  ：kiway
@Date    ：2024/5/1 10:22 
"""

from ..logger import pretty_print

from .BaseEngine import (SwitchEvent, PolyBarEvent)

import bisect
from .BaseEngine import PrOrderEvent
import pandas as pd
from .DataPersist import LatestRecord
import datetime
import numpy as np
from ..logger import pretty_error, pretty_print

time_delta = datetime.timedelta(hours=8)

# 设置交易币种
universe = ['BTCUSDT', 'ETHUSDT']
# 设置基础字段 TODO 千万不要修改这两个list里元素的顺序!!!
future_fields = [
    # future_kline
    'open', 'high', 'low', 'close', 'vwap', 'twap',
    'volume', 'amount', 'count',
    # future_aggTrades
    'act_buy_amount', 'act_buy_volume', 'act_buy_count',
    'act_sell_amount', 'act_sell_volume', 'act_sell_count',
    'BidActMicroVol', 'BidActSmallVol', 'BidActMediumVol', 'BidActLargeVol', 'BidActHugeVol',
    'BidActMicroAmt', 'BidActSmallAmt', 'BidActMediumAmt', 'BidActLargeAmt', 'BidActHugeAmt',
    'BidActMicroCnt', 'BidActSmallCnt', 'BidActMediumCnt', 'BidActLargeCnt', 'BidActHugeCnt',
    'BidActVol', 'BidActAmt', 'BidActCnt',
    'AskActMicroVol', 'AskActSmallVol', 'AskActMediumVol', 'AskActLargeVol', 'AskActHugeVol',
    'AskActMicroAmt', 'AskActSmallAmt', 'AskActMediumAmt', 'AskActLargeAmt', 'AskActHugeAmt',
    'AskActMicroCnt', 'AskActSmallCnt', 'AskActMediumCnt', 'AskActLargeCnt', 'AskActHugeCnt',
    'AskActVol', 'AskActAmt', 'AskActCnt',
    'MicroVol', 'MicroAmt', 'MicroCnt',
    'SmallVol', 'SmallAmt', 'SmallCnt',
    'MediumVol', 'MediumAmt', 'MediumCnt',
    'LargeVol', 'LargeAmt', 'LargeCnt',
    'HugeVol', 'HugeAmt', 'HugeCnt',
    'BidActAmtMed', 'AskActAmtMed',
    'NegCnt', 'NegDis', 'StaCnt', 'MoveDis', 'PosCnt', 'PosDis',
]
spot_fields = [
    # spot_kline
    'spot_open', 'spot_high', 'spot_low', 'spot_close', 'spot_vwap', 'spot_twap',
    'spot_volume', 'spot_amount', 'spot_count',
    'spot_act_buy_volume', 'spot_act_buy_amount', 'spot_act_buy_count',
    'spot_act_sell_volume', 'spot_act_sell_amount', 'spot_act_sell_count',
    # spot_aggTrades
    'spot_BidActMicroVol', 'spot_BidActSmallVol', 'spot_BidActMediumVol', 'spot_BidActLargeVol', 'spot_BidActHugeVol',
    'spot_BidActMicroAmt', 'spot_BidActSmallAmt', 'spot_BidActMediumAmt', 'spot_BidActLargeAmt', 'spot_BidActHugeAmt',
    'spot_BidActMicroCnt', 'spot_BidActSmallCnt', 'spot_BidActMediumCnt', 'spot_BidActLargeCnt', 'spot_BidActHugeCnt',
    'spot_BidActVol', 'spot_BidActAmt', 'spot_BidActCnt',
    'spot_AskActMicroVol', 'spot_AskActSmallVol', 'spot_AskActMediumVol', 'spot_AskActLargeVol', 'spot_AskActHugeVol',
    'spot_AskActMicroAmt', 'spot_AskActSmallAmt', 'spot_AskActMediumAmt', 'spot_AskActLargeAmt', 'spot_AskActHugeAmt',
    'spot_AskActMicroCnt', 'spot_AskActSmallCnt', 'spot_AskActMediumCnt', 'spot_AskActLargeCnt', 'spot_AskActHugeCnt',
    'spot_AskActVol', 'spot_AskActAmt', 'spot_AskActCnt',
    'spot_MicroVol', 'spot_MicroAmt', 'spot_MicroCnt',
    'spot_SmallVol', 'spot_SmallAmt', 'spot_SmallCnt',
    'spot_MediumVol', 'spot_MediumAmt', 'spot_MediumCnt',
    'spot_LargeVol', 'spot_LargeAmt', 'spot_LargeCnt',
    'spot_HugeVol', 'spot_HugeAmt', 'spot_HugeCnt',
    'spot_BidActAmtMed', 'spot_AskActAmtMed',
    'spot_NegCnt', 'spot_NegDis', 'spot_StaCnt', 'spot_MoveDis', 'spot_PosCnt', 'spot_PosDis'
]


# cum_quantile
class CumQuantile:
    def __init__(self):
        self.queue = []
        self.ordered_queue = []
        self.count = 0

    def update(self, data):
        self.ordered_queue.insert(bisect.bisect(self.ordered_queue, data, lo=0), data)
        self.queue.append(data)

    def info(self, q):
        if len(self.queue) > 0:
            idx = int(q * (len(self.queue) - 1))
            if len(self.queue) % 2 == 1:
                return self.ordered_queue[idx]
            else:
                return (self.ordered_queue[idx] + self.ordered_queue[idx + 1]) / 2
        else:
            return np.nan

    def clear(self):
        self.queue.clear()
        self.ordered_queue.clear()
        self.count = 0


class CumAvg:
    def __init__(self) -> None:
        self._avg = np.nan
        self._sum = 0
        self._cnt = 0

    def __call__(self, data):
        if not (np.isinf(data) | np.isnan(data)):
            self._sum += data
            self._cnt += 1.0
            self._avg = self._sum / self._cnt
        return self.info

    @property
    def info(self):
        return self._avg

    def clear(self):
        return self.__init__()


# 计算单个合约基础特征
class CalculateBasicFeature:
    def __init__(self, symbol: str, is_future: bool):
        self.symbol = symbol
        self.is_future = is_future
        # 这里记录单个币种的基础特征,只有1-d,spot这里以future的基础字段名初始化
        # 最后返回时将所有字段加上前缀spot
        if is_future:
            self.basic_data = np.full(shape=len(future_fields), dtype=np.float64, fill_value=0)
        else:
            self.basic_data = np.full(shape=len(spot_fields), dtype=np.float64, fill_value=0)
        # 初始化的open、high、low
        self.basic_data[0] = np.nan
        self.basic_data[1] = -np.inf
        self.basic_data[2] = np.inf
        self.basic_data[3] = np.nan
        self.last_price = np.nan
        self.act_buy_med = CumQuantile()
        self.act_sell_med = CumQuantile()
        self.twap = CumAvg()

    def __call__(self, data: dict):
        # 更新kline
        # 获得最新价格
        new_price = data['price']
        vol = data['amount']  # bianance的amount是我们所理解的成交量(非成交额)
        amt = new_price * vol  # amount需要自己算
        # open = new_price if open == np.nan
        if np.isnan(self.basic_data[0]):
            self.basic_data[0] = new_price
        # high = new_price >= high? new_price:high
        if new_price >= self.basic_data[1]:
            self.basic_data[1] = new_price
        # low = new_price <= low? new_price:low
        if new_price <= self.basic_data[2]:
            self.basic_data[2] = new_price
        # close = new_price
        self.basic_data[3] = new_price
        # twap = avg(new_price)
        self.basic_data[5] = self.twap(new_price)
        # volume
        self.basic_data[6] += vol
        # amount
        self.basic_data[7] += amt
        # count += 1
        self.basic_data[8] += 1.0
        if data['info']['m']:
            # update the median of active sell amount
            self.act_sell_med.update(amt)
            # micro trade
            if amt < 100:
                self.basic_data[33] += vol  # act_micro_sell_vol
                self.basic_data[38] += amt  # act_micro_sell_amt
                self.basic_data[43] += 1.0  # act_micro_sell_cnt
            # small trade
            elif amt < 1000:
                self.basic_data[34] += vol  # act_small_sell_vol
                self.basic_data[39] += amt  # act_small_sell_amt
                self.basic_data[44] += 1.0  # act_small_sell_cnt
            # medium trade
            elif amt < 10000:
                self.basic_data[35] += vol  # act_medium_sell_vol
                self.basic_data[40] += amt  # act_medium_sell_amt
                self.basic_data[45] += 1.0  # act_medium_sell_cnt
            # large trade
            elif amt < 100000:
                self.basic_data[36] += vol  # act_large_sell_vol
                self.basic_data[41] += amt  # act_large_sell_vol
                self.basic_data[46] += 1.0  # act_large_sell_vol
            else:
                self.basic_data[37] += vol  # act_huge_sell_vol
                self.basic_data[42] += amt  # act_huge_sell_vol
                self.basic_data[47] += 1.0  # act_huge_sell_vol
        else:
            # update the median of active buy amount
            self.act_buy_med.update(amt)
            # micro trade
            if amt < 100:
                self.basic_data[15] += vol  # act_micro_buy_vol
                self.basic_data[20] += amt  # act_micro_buy_amt
                self.basic_data[25] += 1.0  # act_micro_buy_cnt
            # small trade
            elif amt < 1000:
                self.basic_data[16] += vol  # act_small_buy_vol
                self.basic_data[21] += amt  # act_small_buy_amt
                self.basic_data[26] += 1.0  # act_small_buy_cnt
            # medium trade
            elif amt < 10000:
                self.basic_data[17] += vol  # act_medium_buy_vol
                self.basic_data[22] += amt  # act_medium_buy_amt
                self.basic_data[27] += 1.0  # act_medium_buy_cnt
            # large trade
            elif amt < 100000:
                self.basic_data[18] += vol  # act_large_buy_vol
                self.basic_data[23] += amt  # act_large_buy_amt
                self.basic_data[28] += 1.0  # act_large_buy_cnt
            else:
                self.basic_data[19] += vol  # act_huge_buy_vol
                self.basic_data[24] += amt  # act_huge_buy_amt
                self.basic_data[29] += 1.0  # act_huge_buy_cnt
        # 记录价格移动距离
        if self.last_price is not None:
            if self.last_price == new_price:
                self.basic_data[70] += 1

            elif self.last_price > new_price:
                self.basic_data[68] += 1
                self.basic_data[69] += (self.last_price - new_price)

            else:
                self.basic_data[71] += 1
                self.basic_data[72] += (new_price - self.last_price)

        self.last_price = new_price

    def switch_time_cut(self):
        # vwap
        if self.basic_data[6] > 0:
            self.basic_data[4] = self.basic_data[7] / self.basic_data[6]
        else:
            self.basic_data[4] = np.nan
        # twap TODO 如果没有成交应该是多少?
        self.basic_data[5] = self.twap.info
        # bid_act_vol
        self.basic_data[30] = self.basic_data[15:20].sum()
        # bid_act_amt
        self.basic_data[31] = self.basic_data[20:25].sum()
        # bid_act_cnt
        self.basic_data[32] = self.basic_data[25:30].sum()
        # ask_act_vol
        self.basic_data[48] = self.basic_data[33:38].sum()
        # ask_act_amt
        self.basic_data[49] = self.basic_data[38:43].sum()
        # ask_act_cnt
        self.basic_data[50] = self.basic_data[43:48].sum()
        # micro_vol, micro_amt, micro_cnt 51 52 53
        self.basic_data[51] = self.basic_data[15] + self.basic_data[33]
        self.basic_data[52] = self.basic_data[20] + self.basic_data[38]
        self.basic_data[53] = self.basic_data[25] + self.basic_data[43]
        # small_vol, small_amt, small_cnt
        self.basic_data[54] = self.basic_data[16] + self.basic_data[34]
        self.basic_data[55] = self.basic_data[21] + self.basic_data[39]
        self.basic_data[56] = self.basic_data[26] + self.basic_data[44]
        # medium_vol, medium_amt, medium_cnt
        self.basic_data[57] = self.basic_data[17] + self.basic_data[35]
        self.basic_data[58] = self.basic_data[22] + self.basic_data[40]
        self.basic_data[59] = self.basic_data[27] + self.basic_data[45]
        # large_vol, large_amt, large_cnt
        self.basic_data[60] = self.basic_data[18] + self.basic_data[36]
        self.basic_data[61] = self.basic_data[23] + self.basic_data[41]
        self.basic_data[62] = self.basic_data[28] + self.basic_data[46]
        # huge_vol, huge_amt, huge_cnt
        self.basic_data[63] = self.basic_data[19] + self.basic_data[37]
        self.basic_data[64] = self.basic_data[24] + self.basic_data[42]
        self.basic_data[65] = self.basic_data[29] + self.basic_data[47]
        # bid_act_amt_med, ask_act_amt_med
        self.basic_data[66] = self.act_buy_med.info(.5)
        self.basic_data[67] = self.act_sell_med.info(.5)
        # reset basic fields
        res_df = self.basic_data + 0
        self.basic_data.fill(0)
        # open
        self.basic_data[0] = np.nan
        # high
        self.basic_data[1] = -np.inf
        # low
        self.basic_data[2] = np.inf
        # close
        self.basic_data[3] = np.nan
        # reset buy_med, sell_med, twap
        self.act_buy_med = CumQuantile()
        self.act_sell_med = CumQuantile()
        self.twap = CumAvg()
        return res_df


# 计算所有基础特征
class BaiscFeature:
    def __init__(self, universe):
        if universe is None:
            # raise Exception("universe must be unique")
            return

        # 最终返回的dataframe的
        # columns
        self.universe = universe
        # index
        self.field_names = future_fields + spot_fields
        # 初始化基础特征矩阵
        self.data = np.full(
            shape=(len(future_fields) + len(spot_fields), 1),
            dtype=np.float64,
            fill_value=np.nan,
        )
        self.time_cut = None
        # 初始化计算引擎
        self.future_dict = CalculateBasicFeature(symbol=universe, is_future=True)

        self.spot_dict = CalculateBasicFeature(symbol=universe, is_future=False)

        self.res_data = None

    def __call__(self, symbol: str, data: dict, is_future: bool, tms):
        # pretty_print(f"{tms} -- update aggTrades: {symbol}, {'future' if is_future else 'spot'}")
        # symbol = data['info']['s']
        # 更新期货基础特征
        if is_future:

            self.future_dict(data)

        # 更新现货基础特征
        else:
            self.spot_dict(data)

    def switch_time_cut(self):
        future_idx = self.data.shape[0] // 2
        # for i, symbol in enumerate(self.universe):
        self.data[:future_idx, 0] = self.future_dict.switch_time_cut()
        self.data[future_idx:, 0] = self.spot_dict.switch_time_cut()
        res_df = pd.DataFrame(self.data, index=self.field_names, columns=[self.universe])
        return res_df


class Strategy__:
    flag = True

    def __init__(self, _id=None):

        self.position_info_df = None
        self.balance_info_df = None
        self.order_status_df = None

        self._id = _id
        self.main_engine = None
        if Strategy__.flag:
            self.factor_init()
            Strategy__.flag = False

        self.basic_feature_engine = BaiscFeature(universe=self._id)
        self.is_future = self._id.endswith(":USDT")
        pretty_print(f"{self._id} 子策略启动")

        # self.time_cut = pd.to_datetime(datetime.datetime.now()).floor('min')

    def get_time_cut_backward(self):
        return self.record.time_cut_backward()

    def get_time_cut_forward(self):
        return self.record.time_cut_forward()

    def record_update(self, sid, value):
        self.record.update(sid, value)

    #
    @property
    def record(self):
        return LatestRecord()

    @classmethod
    def factor_init(cls):
        pass

    def set_main_engine(self, main_engine):
        self.main_engine = main_engine

    def _set_id(self, value: str):
        self._id = value
        pretty_print(f"子进程-{value} 初始化")

    def _generate_bar_from_aggtrade(self, symbol, aggTrade):
        """从aggtrade合成bar"""

        # TODO 这里时间戳还是没调整过来,暂时得手动+8h才是正常时间
        time_cut = pd.to_datetime(aggTrade[0]['timestamp'], unit='ms').floor('min') + time_delta

        if str(time_cut) > self.get_time_cut_forward():
            self.record_update("time_cut_forward", str(time_cut))

        self.basic_feature_engine(
            symbol=self._id
            , data=aggTrade[0]  # 第一条为最新的归集交易信息
            , is_future=self.is_future
            , tms=aggTrade[0]["timestamp"]
        )

        # 判断是否更新截面,如果跨越了新的一分钟,那么更新该截面
        # print(f"time_cut_forward is {self.get_time_cut_forward()}, time_cut_backward is {self.get_time_cut_backward()}")
        if self.get_time_cut_forward() > self.get_time_cut_backward():
            self.record_update("time_cut_backward", self.get_time_cut_forward())

            self._put_event(SwitchEvent())
            pretty_print("推送信号")

    def switch_time_cut(self):
        res_df = self.basic_feature_engine.switch_time_cut()
        # TODO 以下写往下游推送的逻辑
        self._put_event(PolyBarEvent(res_df))

    def place_order_from_all_bar(self, order_data: pd.DataFrame):
        """
        适用于截面策略下单
        :param order_data:  pd.Series[index.symbol, value.factor_value
        :return:
        """
        exchange_id = self._id.split("--")[0]
        assert order_data["value"].abs().sum() - 1 < 0.00001, "因子没有归一化，请纠正"
        self._put_event(PrOrderEvent("ALL", order_data, exchange_id))

    def place_order(self, symbol, order_data):
        """对单个标的进行下单"""
        # symbol = self._id.split("--")[0]
        exchange_id = self._id.split("--")[0]
        self._put_event(PrOrderEvent(symbol, order_data, exchange_id))

    def on_bar(self, symbol, bars):
        pass

    def on_tick(self, symbol, ticks):
        pass

    def on_orderbook(self, symbol, orderbooks):
        pass

    def on_all_bar(self, bars):
        pass

    def on_aggTrade(self, symbol, aggTrade):
        pass

    def _put_event(self, event):
        self.main_engine.put_event(event)
